package tech.ch.udf

import com.demo.utils.IPUtils
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.sql.SparkSession
import scala.collection.JavaConverters._

class SparkIpUDF(spark: SparkSession) {

  def registerCommon(): Unit = {
    spark.udf.register("ipv4_to_num", (ipv4: String) => {
      IPUtils.ipv4ToNum(ipv4)
    })

    spark.udf.register("num_to_ipv4", (num: Long) => {
      IPUtils.num2Ipv4(num)
    })

    spark.udf.register("ipv4_short", (ipv4: String) => {
      IPUtils.convertIpv4ToShort(ipv4)
    })

    spark.udf.register("ipv4_full", (ipv4: String) => {
      IPUtils.convertIpv4ToFull(ipv4)
    })
  }



  /**
    * 1. 广播变量的值一定要在UDF函数的定义中获取
    * 2. 不能对广播变量的值做转换，比如字符串的截取
    * @param ipLibBD
    */
  def registerIpSearch(ipLibBD:Broadcast[Array[(String, String, String)]]): Unit = {
    spark.udf.register("ip_search", (ip: String) => {
      val ipLibArr = ipLibBD.value
      // 生成开始ip和结束ip的数组。 [[1.1.1.1, 1.1.1.2], [2.2.2.1, 2.2.2.10]]
      var ipRangeArr = ipLibArr.map(x => List(x._1, x._2).toArray)
      val index = IPUtils.ipBinarySearch(ipRangeArr, IPUtils.convertIpv4ToFull(ip))
      if(index >= 0) {
        ipLibArr(index)._3
      } else {
        "unknown"
      }
    })
  }


}
